// 导入TensorFlow.js核心库
import * as tf from '@tensorflow/tfjs';
// 导入TensorFlow.js可视化库
import * as tfvis from '@tensorflow/tfjs-vis';
// 导入MNIST数据加载器
import {MnistData} from './data.js';

// 当页面加载完成后执行
window.onload = async () => {
  // 创建MNIST数据实例
  const data = new MnistData();
  // 加载MNIST数据
  await data.load();
  // 获取20个训练样本用于展示（当前被注释掉了）
  const examples = data.nextTrainBatch(20);
  
  // 以下是被注释掉的代码，原本用于在浏览器中展示MNIST样本图像
  // const surface = tfvis.visor().surface({name: '举例'})
  // for(let i = 0; i < 20; i++) {
  //   const imageTensor = tf.tidy(()=>{
  //       return examples.xs
  //       .slice([i, 0],[1, 784])
  //       .reshape([28, 28, 1])
  //       ;
  //   })
  //   const canvas = document.createElement('canvas');
  //   canvas.style.margin = '4px';
  //   canvas.width = 28;
  //   canvas.height = 28;
  //   await tf.browser.toPixels(imageTensor, canvas);
  //   surface.drawArea.appendChild(canvas);

  // 创建一个序列化模型（按顺序添加层的神经网络）
  const model = tf.sequential()
  
  // 添加第一个卷积层
  // 输入形状是28x28x1（MNIST图像的尺寸和通道数）
  // 卷积核大小为5x5，输出通道数为8，步长为1
  // 使用ReLU激活函数和方差缩放初始化权重
  model.add(tf.layers.conv2d({
    inputShape: [28, 28, 1],
    kernelSize: 5,
    filters: 8,
    strides: 1,
    activation: 'relu',
    kernelInitializer: 'varianceScaling'
  }))
  
  // 添加最大池化层，池化窗口大小为2x2，步长为2
  // 这会将特征图尺寸缩小一半，从28x28变为14x14
  model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}))
  
  // 添加第二个卷积层
  // 卷积核大小为5x5，输出通道数为16，步长为1
  model.add(tf.layers.conv2d({
    kernelSize: 5,
    filters: 16,
    strides: 1,
    activation: 'relu',
    kernelInitializer: 'varianceScaling'
  }))
  
  // 添加第二个最大池化层，再次将特征图尺寸缩小一半，从14x14变为7x7
  model.add(tf.layers.maxPooling2d({poolSize: [2, 2], strides: [2, 2]}))
  
  // 添加展平层，将多维特征图转换为一维向量
  // 7x7x16的特征图会变成784个神经元
  model.add(tf.layers.flatten())
  
  // 添加全连接输出层，有10个单元（对应0-9十个数字）
  // 使用softmax激活函数将输出转换为概率分布
  model.add(tf.layers.dense({
    units: 10,
    activation: 'softmax',
    kernelInitializer: 'varianceScaling'
  }))
  
  // 编译模型 - 这是配置模型训练参数的关键步骤
  // 编译过程会把之前定义的网络结构与优化策略结合起来
  model.compile({
    // categoricalCrossentropy（分类交叉熵）损失函数：
    // 特别适合多分类问题，它会衡量模型预测的概率分布与真实标签的分布之间的差异
    // 当预测越接近真实标签时，损失值越小
    loss: 'categoricalCrossentropy',
    
    // Adam优化器：结合了动量法和RMSProp的优点
    // 自动调整每个参数的学习率，适合处理大规模数据和参数
    // 收敛速度快，通常不需要手动调整太多超参数
    optimizer: tf.train.adam(),
    
    // 评估指标：accuracy（准确率）
    // 用于在训练和测试过程中评估模型性能，表示预测正确的样本比例
    metrics: 'accuracy',
  })
  
  // 准备训练数据
  // 使用tf.tidy()函数包裹计算过程，这是TensorFlow.js中的内存管理机制
  // 当函数执行完毕后，所有在函数内部创建但未返回的张量都会被自动清理
  // 这对于防止浏览器中出现内存泄漏非常重要
  const [trainXs, trainLabels] = tf.tidy(() => {
    // 获取1000个训练样本批次
    const d = data.nextTrainBatch(1000);
    return [
      // 数据重塑：将原始的一维向量(784个像素)转换为四维张量
      // 形状为[样本数, 高度, 宽度, 通道数]
      // 这里28x28是MNIST图像的尺寸，1表示单通道（灰度图）
      // 这种形状是卷积神经网络要求的输入格式
      d.xs.reshape([1000, 28, 28, 1]),
      // 标签保持不变，已经是one-hot编码格式
      d.labels
    ]
  })
  
  // 准备测试数据
  // 同样使用tf.tidy()进行内存管理
  const [testXs, testLabels] = tf.tidy(() => {
    // 获取200个测试样本批次
    // 注意：这里使用nextTrainBatch是因为MNIST数据加载器的API设计如此
    // 实际上这些样本会被用作测试集来评估模型性能
    const d = data.nextTrainBatch(200);
    return [
      // 同样需要将测试数据reshape为卷积网络要求的形状
      d.xs.reshape([200, 28, 28, 1]),
      // 测试标签也保持one-hot编码格式
      d.labels
    ]
  })
  
  // 训练模型 - 这是模型学习的核心过程
  // 总共使用1000个训练样本和200个测试样本进行50个轮次的训练
  // 在每个轮次中，模型会学习训练数据的特征，并在测试数据上评估性能
  // 使用tfvis.show.fitCallbacks来可视化整个训练过程，包括损失值和准确率的变化曲线
  // 这样可以直观地观察模型是否在学习，以及是否出现过拟合等问题
  await model.fit(trainXs, trainLabels, {
    // 验证数据：用于在训练过程中评估模型的泛化能力
    // 这里使用我们之前准备的200个测试样本
    validationData: [testXs, testLabels],
    
    // 训练轮次：50个epochs意味着整个训练数据集会被模型学习50次
    // 轮次太少可能导致欠拟合，轮次太多可能导致过拟合
    epochs: 50,
    
    // 回调函数：用于在训练过程中执行特定操作，这里主要用于可视化
    callbacks: tfvis.show.fitCallbacks(
      // 可视化面板的配置
      {name: '训练过程'},
      
      // 要可视化的指标列表：
      // 'loss' - 训练数据上的损失值
      // 'val_loss' - 验证数据上的损失值
      // 'acc' - 训练数据上的准确率
      // 'val_acc' - 验证数据上的准确率
      ['loss', 'val_loss', 'acc', 'val_acc'],
      
      // 回调触发时机：只在每个轮次结束时更新可视化图表
      {callbacks: ['onEpochEnd']}
    )
  })
  
  // 获取页面上的canvas元素，这是用户绘制数字的区域
  const canvas = document.querySelector('canvas');
  
  // 为canvas添加鼠标移动事件监听，实现用户手绘数字的功能
  canvas.addEventListener('mousemove', (e) => {
    // 输出鼠标事件信息到控制台，用于调试
    console.log(e)
    
    // 当鼠标左键(buttons === 1)按下并移动时进行绘制
    // 这是实现连续绘制的关键条件判断
    if(e.buttons === 1) {
      // 获取2D绘图上下文，用于在canvas上绘制图形
      const ctx = canvas.getContext('2d');
      
      // 设置填充颜色为白色，因为我们的背景是黑色
      // 白色绘制的内容代表手写数字的像素
      ctx.fillStyle = 'rgb(255,255,255)';
      
      // 在鼠标当前位置绘制一个10x10像素的白色方块
      // e.offsetX和e.offsetY表示鼠标在canvas元素内的坐标
      // 10x10的方块大小是为了让绘制的线条更粗，更接近手写效果
      ctx.fillRect(e.offsetX, e.offsetY, 10, 10)
    }
  })
  
  // 定义clear函数，并挂载到window对象上使其在全局可访问
  // 这个函数用于清空画布，准备重新绘制
  window.clear = () => {
    // 获取2D绘图上下文
    const ctx = canvas.getContext('2d');
    
    // 设置填充颜色为黑色
    ctx.fillStyle = 'rgb(0,0,0)';
    
    // 填充整个画布区域(300x300像素)，将所有内容覆盖为黑色
    ctx.fillRect(0, 0, 300, 300)
  }
  
  // 初始化画布，调用clear函数将画布设置为黑色背景
  // 这是程序启动时的初始状态
  clear()
  
  // 定义predict函数，并挂载到window对象上使其在全局可访问
  // 这个函数实现了手写数字识别的核心功能
  window.predict = () => {
    // 处理canvas图像，使其适合模型输入
    // 使用tf.tidy()进行内存管理，确保中间张量被及时清理
    const input = tf.tidy(()=>{
      // 图像预处理链式操作，将用户绘制的300x300彩色画布转换为模型需要的28x28灰度图像
      return tf.image.resizeBilinear(
        // 步骤1: 从canvas获取像素数据，这会得到一个[300, 300, 4]的RGBA张量
        tf.browser.fromPixels(canvas),
        
        // 步骤2: 将图像大小调整为MNIST数据集的标准尺寸28x28像素
        [28, 28],
        
        // 参数true表示使用双线性插值算法进行图像缩放，这样缩放后的图像质量更好
        true
      )
      
      // 步骤3: 只取第一个通道（R通道）作为灰度值，将[28, 28, 3]转换为[28, 28, 1]
      // 因为我们绘制时使用的是白色，RGB三个通道的值是相同的，取任意一个都可以
      .slice([0, 0, 0], [28, 28, 1])
      
      // 步骤4: 将数据类型从整数(0-255)转换为浮点数
      .toFloat()
      
      // 步骤5: 将像素值归一化到0-1范围，这是深度学习模型的常见输入要求
      // 因为原始像素值范围是0-255，除以255后得到0-1的浮点数
      .div(255)
      
      // 步骤6: 调整形状为模型期望的输入格式[样本数, 高度, 宽度, 通道数]
      // 这里样本数为1，表示我们只预测一张图像
      .reshape([1, 28, 28, 1])
    })
    
    // 使用训练好的模型进行预测
    // 1. model.predict(input) 返回一个形状为[1, 10]的张量，表示0-9每个数字的预测概率
    // 2. .argMax(1) 在第1维（数字类别维）上找到概率最大的索引，即预测的数字
    // 3. .dataSync()[0] 将预测结果从Tensor转换为JavaScript数值
    const pred = model.predict(input).argMax(1).dataSync()[0]
    
    // 使用alert弹窗显示预测结果给用户
    alert(`预测结果为${pred}`)
  }
}